/*
 * Copyright 2012 David Crosson
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package fr.janalyse.jmx

import scala.util.{ Properties => SP }
import java.io.File.{ separator => FS, pathSeparator => PS }
import scala.reflect.BeanProperty
import java.lang.management.ManagementFactory
import javax.management.ObjectInstance
import javax.management.ObjectName
import javax.management.remote.rmi.RMIConnection
import javax.management.remote.rmi.RMIServer
import javax.management.remote.rmi.RMIServerImpl_Stub
import javax.management.openmbean.{ CompositeData, CompositeDataSupport }
import javax.management.openmbean.{ TabularData, TabularDataSupport }
import java.rmi.MarshalledObject
import java.rmi.UnmarshalException
import javax.management.openmbean.CompositeData
import scala.collection.JavaConversions._
import javax.management.MBeanAttributeInfo
import java.rmi.registry.LocateRegistry
import javax.management.remote.JMXConnector
import javax.management.remote.JMXServiceURL
import javax.management.remote.{JMXConnectorServerFactory, JMXConnectorFactory}
import java.rmi.server.RMIClientSocketFactory
import javax.management.remote.rmi.RMIJRMPServerImpl
import java.rmi.server.RMIServerSocketFactory
import javax.net.ServerSocketFactory
import java.net.ServerSocket
import java.rmi.server.RMISocketFactory
import javax.management.RuntimeMBeanException
import scala.math.ScalaNumber

case class CompositeDataWrapper[CDT <: CompositeData](cd: CDT) {
  lazy val content: Map[String, Object] = {
    val keys = cd.getCompositeType().keySet()
    val tuples = keys.toList map { k => (k -> cd.get(k)) }
    tuples.toMap
  }
  def toWrapper() = this
  //def get(key:String):Option[Object] = content.get(key)
  def getString(key: String): Option[String] = content.get(key) map { _.toString }
  def get[A](key: String): Option[A] = content.get(key) map { _.asInstanceOf[A] }

  // TODO Improve number support and add more typed getter for basic types
  def getNumber[N >: Number](key: String): Option[N] = {
    for (raw <- content.get(key)) yield {
      raw match {
        case e: java.lang.Byte => e.toDouble
        case e: java.lang.Short => e.toDouble
        case e: java.lang.Integer => e.toDouble
        case e: java.lang.Long => e.toDouble
        case e: java.lang.Float => e.toDouble
        case e: java.lang.Double => e
        case e: java.lang.String => e.toDouble
      }
    }
  }
}

case class TabularDataWrapper[TD <: TabularData](tabularData: TD) {
  lazy val content: Map[String, Map[String, Object]] = {
    val indexNamesInCell = tabularData.getTabularType().getIndexNames()
    val tuples = tabularData.values collect {
      case v: CompositeDataSupport =>
        val cellContent = v.content
        val key = indexNamesInCell map { cellContent.get(_).get } mkString ("-")
        val submap = cellContent filterNot { case (k, _) => indexNamesInCell contains k }

        key -> submap

      case (k: java.util.List[String], v: CompositeData) =>
        val cellContent = v.content
        val key = indexNamesInCell map { cellContent.get(_).get } mkString ("-")
        val submap = cellContent filterNot { case (k, _) => indexNamesInCell contains k }

        key -> submap
    }
    tuples.toMap
  }
  def toWrapper() = this
  def get(key: String): Option[Map[String, Object]] = content.get(key)
  def get(key: String, subkey: String): Option[Object] = content.get(key) flatMap { _.get(subkey) }
  def getString(key: String, subkey: String): Option[String] = get(key, subkey) map { _.toString }
}

object RichMBean {
  def apply(name: String)(implicit jmx: JMX) = new RichMBean(name)
}

case class RichMBean(objectName: ObjectName)(implicit jmx: JMX) {
  val name = objectName.toString()
  val keys = mapAsScalaMap(objectName.getKeyPropertyList())

  private lazy val objectInstance = jmx.rmiConnection.getObjectInstance(objectName, null) match {
    case null => throw new RuntimeException("Couldn't find MBean")
    case mbean => mbean
  }

  lazy val attributesMap:Map[String, RichAttribute] = {
     try {
	    val mbeanInfo = jmx.rmiConnection.getMBeanInfo(objectName, null)
	    val attrsInfo: List[MBeanAttributeInfo] = mbeanInfo.getAttributes().toList
	    attrsInfo map { ai =>
	      val n = ai.getName()
	      val d = if (ai.getDescription() == ai.getName) None else Some(ai.getDescription())
	      val a: RichAttribute = ai.getType() match {
	        case "java.lang.Boolean" | "boolean" | "java.lang.boolean" => RichBooleanAttribute(n, d)
	        case "java.lang.Byte" | "byte" => RichByteAttribute(n, d)
	        case "java.lang.Short" | "short" => RichShortAttribute(n, d)
	        case "java.lang.Integer" | "int" => RichIntAttribute(n, d)
	        case "java.lang.Long" | "long" => RichLongAttribute(n, d)
	        case "java.lang.Float" | "float" => RichFloatAttribute(n, d)
	        case "java.lang.Double" | "double" => RichDoubleAttribute(n, d)
	        case "java.lang.String" | "String" => RichStringAttribute(n, d)
	        //case "java.util.List" =>
	        //case "java.util.Properties" =>
	        case "[Ljava.lang.String;" => RichStringArrayAttribute(n, d)
	        //case "javax.management.openmbean.TabularData"      =>
	        case "javax.management.openmbean.CompositeData" | "javax.management.openmbean.CompositeDataSupport" => RichCompositeDataAttribute(n, d)
	        //case "[Ljavax.management.openmbean.CompositeData;" =>
	        //case "[Ljavax.management.ObjectName;" => 
	        case x =>
	          //println("Warning: Not supported jmx attribute value type %s for %s mbean %s".format(x, n, name))
	          //println(x)
	          RichGenericAttribute(ai.getName(), d)
	      }
	      n -> a
	    } toMap
	    //def apply[A](attrname:String) = jmx.rmiConnection.getAttribute(objectName,attrname,null).asInstanceOf[A]
    } catch {
      case e =>
        println("Couln't build attributes map for MBean %s (%s)".format(name, e.getMessage))
        Map.empty[String,RichAttribute]
    }
  }

  def get[A](name: String): Option[A] = {
    val mbeanInfo = jmx.rmiConnection.getMBeanInfo(objectName, null)
    val attrsInfo: List[javax.management.MBeanAttributeInfo] = mbeanInfo.getAttributes().toList
    (attrsInfo find { attrInfo => attrInfo.getName == name }) match {
      case None => None
      case Some(attrInfo) => apply[A](name) match {
        case x if (x == null) => None
        case x => Some(x)
      }
    }
  }
  def apply[A](attrname: String) = jmx.rmiConnection.getAttribute(objectName, attrname, null).asInstanceOf[A]
  def set[A](attrname: String, value: A) {
    val attribute = new javax.management.Attribute(attrname, value)
    jmx.rmiConnection.setAttribute(objectName, attribute, null)
  }
  def apply[A](attrnames: String*): List[A] = attrnames.toList map { apply[A](_) }
  def call[A](operation: String, args: Any*): Option[A] = {
    val signature = args map { // http://home.pacifier.com/~mmead/jni/cs510ajp/index.html  /   http://en.wikipedia.org/wiki/Java_Native_Interface#Mapping_types
      case _: Int => "int"
      case _: Short => "short"
      case _: Long => "long"
      case _: Float => "float"
      case _: Double => "double"
      case _: Boolean => "boolean"
      case _: Byte => "byte"
      case _: Char => "char"
      case _: String => "String"
      case _: Array[Int] => "[I"
      case _: Array[Short] => "[S"
      case _: Array[Long] => "[J"
      case _: Array[Float] => "[F"
      case _: Array[Double] => "[D"
      case _: Array[Boolean] => "[Z"
      case _: Array[Byte] => "[B"
      case _: Array[Char] => "[C"
      case _: Array[String] => "[Ljava/lang/String"
    }
    val args2give = args.toArray
    jmx.rmiConnection.invoke(objectName, operation, args2give, signature.toArray, null) match {
      case null => None
      case result => Some(result.asInstanceOf[A])
    }
  }

  private def genericGetter[T, R <: RichAttribute](attr: R, getter: (Object) => T): Option[T] = {
    try {
      jmx.rmiConnection.getAttribute(objectName, attr.name, null) match {
        case null => None
        case x => Some(getter(x))
      }
    } catch {
      case e: RuntimeMBeanException if e.getCause().isInstanceOf[UnsupportedOperationException] => None
      case e: UnmarshalException => None
      case e: java.rmi.ConnectException => throw e
      case x =>
        //println("Warning: Error while getting value for attribute %s mbean %s (%s)".format(attr.name, name, x))
        None
    }
  }

  def getString(attr: RichAttribute): Option[String] = genericGetter(attr, attr.asString)

  def getDouble(attr: RichNumberAttribute): Option[Double] = genericGetter(attr, attr.asDouble)

  def getLong(attr: RichNumberAttribute): Option[Long] = genericGetter(attr, attr.asLong)

  def getInt(attr: RichNumberAttribute): Option[Int] = genericGetter(attr, attr.asInt)

  def getComposite(attr: RichCompositeDataAttribute): Option[Map[String, Object]] = genericGetter(attr, attr.asMap)

  def getNumberComposite[N >: Number](attr: RichCompositeDataAttribute): Option[Map[String, N]] = genericGetter(attr, attr.asNumberMap[N])

  def getString(attrname: String): Option[String] = attributesMap.get(attrname).flatMap(getString(_))

  def getDouble(attrname: String): Option[Double] = attributesMap.get(attrname).collect({ case x: RichNumberAttribute => x }).flatMap(getDouble(_))

  def getLong(attrname: String): Option[Long] = attributesMap.get(attrname).collect({ case x: RichNumberAttribute => x }).flatMap(getLong(_))

  def getInt(attrname: String): Option[Int] = attributesMap.get(attrname).collect({ case x: RichNumberAttribute => x }).flatMap(getInt(_))
  
  def getComposite(attrname: String): Option[Map[String, Object]] = attributesMap.get(attrname).collect({ case x: RichCompositeDataAttribute => x }).flatMap(getComposite(_))

  def getNumberComposite[N >: Number](attrname: String): Option[Map[String, N]] = attributesMap.get(attrname).collect({ case x: RichCompositeDataAttribute => x }).flatMap(getNumberComposite(_))

  def attributes(): Iterable[RichAttribute] = attributesMap.values
}

/*
import fr.janalyse.jmx._
val jmx=JMX()
val os=jmx.os.get
val mem=jmx.memory.get
val th=jmx.threading.get
 */

trait RichAttribute {
  val name: String
  val desc: Option[String]
  def asString(ob: Object): String = ob.toString() // TODO BAD
}

trait RichArrayAttribute extends RichAttribute {
}

trait RichNumberAttribute extends RichAttribute {
  def asDouble(ob: Object): Double = ob match { // TODO BAD
    case e: java.lang.Byte => e.toDouble
    case e: java.lang.Short => e.toDouble
    case e: java.lang.Integer => e.toDouble
    case e: java.lang.Long => e.toDouble
    case e: java.lang.Float => e.toDouble
    case e: java.lang.Double => e
    case e: java.lang.String => e.toDouble
  }
  def asLong(ob: Object): Long = ob match { // TODO BAD
    case e: java.lang.Byte => e.toLong
    case e: java.lang.Short => e.toLong
    case e: java.lang.Integer => e.toLong
    case e: java.lang.Long => e
    case e: java.lang.Float => e.toLong
    case e: java.lang.Double => e.toLong
    case e: java.lang.String => e.toLong
  }
  def asInt(ob: Object): Int = ob match { // TODO BAD
    case e: java.lang.Byte => e.toInt
    case e: java.lang.Short => e.toInt
    case e: java.lang.Integer => e
    case e: java.lang.Long => e.toInt
    case e: java.lang.Float => e.toInt
    case e: java.lang.Double => e.toInt
    case e: java.lang.String => e.toInt
  }

}

case class RichCompositeDataAttribute(name: String, desc: Option[String] = None) extends RichAttribute {
  override def asString(ob: Object): String = {
    val cd = ob.asInstanceOf[CompositeData] // TODO BAD
    cd.toString()
  }
  def asMap(ob: Object) = ob.asInstanceOf[CompositeData].content // TODO BAD
  def asNumberMap[N >: Number](ob: Object): Map[String, N] = {
    val cd = ob.asInstanceOf[CompositeData] // TODO BAD
    cd.content flatMap {
      case (name, value) =>
        value match {
          case x: java.lang.Long => Some(name -> new runtime.RichDouble(x.toDouble))
          case x: java.lang.Double => Some(name -> new runtime.RichDouble(x.toDouble))
          case x: java.lang.Integer => Some(name -> new runtime.RichDouble(x.toDouble))
        }
    }
  }
}

case class RichStringArrayAttribute(name: String, desc: Option[String] = None) extends RichArrayAttribute {
  override def asString(ob: Object): String = { // TODO BAD
    ob.asInstanceOf[Array[String]].toList.mkString(", ")
  }
}

case class RichBooleanAttribute(name: String, desc: Option[String] = None) extends RichAttribute

case class RichStringAttribute(name: String, desc: Option[String] = None) extends RichAttribute

case class RichGenericAttribute(name: String, desc: Option[String] = None) extends RichAttribute

case class RichByteAttribute(name: String, desc: Option[String] = None) extends RichNumberAttribute
case class RichShortAttribute(name: String, desc: Option[String] = None) extends RichNumberAttribute
case class RichIntAttribute(name: String, desc: Option[String] = None) extends RichNumberAttribute
case class RichLongAttribute(name: String, desc: Option[String] = None) extends RichNumberAttribute
case class RichDoubleAttribute(name: String, desc: Option[String] = None) extends RichNumberAttribute
case class RichFloatAttribute(name: String, desc: Option[String] = None) extends RichNumberAttribute

case class Credentials(username: String = "", password: String = "")

case class JMXOptions(
  host: String = "localhost",
  port: Int = 1099,
  name: Option[String] = None,
  username: Option[String] = None,
  password: Option[String] = None,
  connectTimeout: Long = 30000,
  retryCount: Int = 5,
  retryDelay: Long = 2000) {
  val credentials: Option[Credentials] = (username, password) match {
    case (None, None) => None
    case (None, _) => None
    case (Some(u), Some(p)) => Some(Credentials(u, p))
  }
}

trait JMXAutoClose {
  def usingJMX[T <: { def close() }, R](resource: T)(block: T => R) = {
    try block(resource)
    finally resource.close
  }
}

object JMX extends JMXAutoClose {
  def register(ob: Object, obname: ObjectName) = ManagementFactory.getPlatformMBeanServer.registerMBean(ob, obname)

  def once[R](
    host: String,
    port: Int,
    username: Option[String] = None,
    password: Option[String] = None)(block: JMX => R): R = {
    once[R](JMXOptions(host = host, port = port, username = username, password = password))(block)
  }

  def once[R]()(block: JMX => R) = once[R](None)(block)
  def once[R](options: JMXOptions)(block: JMX => R) = once[R](Some(options))(block)
  def once[R](someOptions: Option[JMXOptions])(block: JMX => R) = usingJMX(new JMX(someOptions)) { block(_) }

  def apply(options: JMXOptions) = new JMX(Some(options))
  def apply(options: Option[JMXOptions]) = new JMX(options)
  def apply(host: String, port: Int) = new JMX(Some(JMXOptions(host = host, port = port)))
  def apply() = new JMX()

  def checkServiceURL(urlstr:String):String = {
	  val url  = new JMXServiceURL(urlstr)
	  var jmxc:Option[JMXConnector] = None
	  try {
	    jmxc = Some(JMXConnectorFactory.connect(url))
 	    val mbsc = jmxc.get.getMBeanServerConnection()
  	    val domains = mbsc.getDomains()
	  } finally { jmxc.foreach(_.close) }
	  urlstr
  }
  
  def findMBeanServers(host: String, port: Int, credentials: Option[Credentials]): List[RMIConnection] = {
    def searchMissingClass(e: Throwable): String = {
      e match {
        case null => "?"
        case ex: ClassNotFoundException => ex.getMessage().trim().split("\\s+")(0)
        case _ => searchMissingClass(e.getCause)
      }
    }
    try {
      val lr = java.rmi.registry.LocateRegistry.getRegistry(host, port)
      //filterNot => some jonas MBeans causes registry data corruption if lookuped then all others call are in exception
      val entries = lr.list.toList
        .filterNot(_ contains "TMFactory")
        .filterNot(_ contains "RMI_SERVER_RPC")
      entries map { entry =>
        try {
          val credentialsArray: Array[String] = credentials match {
            case None => Array()
            case Some(Credentials(username, password)) => Array(username, password)
          }
          lr.lookup(entry) match {
            case stub: RMIServerImpl_Stub => Some(stub.newClient(credentialsArray))
            case rsi: RMIServer => Some(rsi.newClient(credentialsArray))
            case r @ _ => println("Do not know how to manage %s class %s".format(entry, r.getClass)); None
          }
        } catch {
          case e: UnmarshalException => /*println("%s - Missing = %s".format(entry, searchMissingClass(e)));*/ None
        }
      } flatten
    }
  }
}



class JMX(val options: Option[JMXOptions] = None) {
  //def jmx=this
  lazy val rmiConnection = options map { o =>
    JMX.findMBeanServers(o.host, o.port, o.credentials).head
  } getOrElse {
    // By default - self connect...
    // TODO : BAD BAD BAD - Find a better way...
    val csf = RMISocketFactory.getDefaultSocketFactory()
    var chosenPort: Option[Int] = None
    val ssf = new RMIServerSocketFactory {
      def createServerSocket(port: Int): ServerSocket = {
        val ssock = ServerSocketFactory.getDefault().createServerSocket(port)
        chosenPort = Some(ssock.getLocalPort())
        ssock
      }
    }
    LocateRegistry.createRegistry(0, csf, ssf)
    val url = new JMXServiceURL("service:jmx:rmi:///jndi/rmi://localhost:%d/jmxapitestrmi".format(chosenPort.get))
    val mbs = ManagementFactory.getPlatformMBeanServer()
    val cs = JMXConnectorServerFactory.newJMXConnectorServer(url, null, mbs)
    cs.start()
    Thread.sleep(500)
    JMX.findMBeanServers("localhost", chosenPort.get, None).head
  }

  def close() = {
    try {
      rmiConnection.close
    } catch {
      case e => println("Exception while closing rmi connection, let's ignore and continue...")
    }
  }

  def apply(name: String) = RichMBean(name)(this)
  def domains = rmiConnection.getDomains(null).toList
  def get(name: String) = if (rmiConnection.queryNames(name, null, null).size > 0) Some(apply(name)) else None
  //
  def mbeans(query: String) = {
    rmiConnection.queryNames(null, string2objectName(query), null).toList map { RichMBean(_)(this) }
  }
  def mbeans() = {
    rmiConnection.queryNames(null, null, null).toList map { RichMBean(_)(this) }
  }
  def os() = get("java.lang:type=OperatingSystem")
  def runtime() = get("java.lang:type=Runtime")
  def threading() = get("java.lang:type=Threading")
  def memory() = get("java.lang:type=Memory")

  def systemProperties() = {
    var sysprops = Map.empty[String, String]
    for (
      rt <- runtime;
      props: TabularDataSupport <- rt.get("SystemProperties");
      (key, values) <- props.content;
      value <- values.get("value")
    ) { sysprops += key -> value.toString }
    sysprops
  }

}

